import torch
import torch.nn as nn
import torchvision
import torch.nn.functional as F

class GTF(nn.Module):
    """
    Gradient Transfer Fusion: 把可见光梯度迁移到红外，然后 RGB 输出
    """

    def __init__(self):
        super().__init__()

    def forward(self, vis, ir):
        """
        vis: (B,3,H,W)  0-1
        ir : (B,1,H,W)  0-1
        return: (B,3,H,W)
        """
        # 计算可见光梯度幅度
        grad_x = F.conv2d(
            vis,  # 3 通道同时卷
            torch.tensor(
                [[[[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]]],
                dtype=vis.dtype,
                device=vis.device,
            ).expand(3, 1, 3, 3),
            groups=3,
            padding=1,
        )
        grad_y = F.conv2d(
            vis,
            torch.tensor(
                [[[[-1, -2, -1], [0, 0, 0], [1, 2, 1]]]],
                dtype=vis.dtype,
                device=vis.device,
            ).expand(3, 1, 3, 3),
            groups=3,
            padding=1,
        )
        grad_mag = torch.sqrt(grad_x.pow(2) + grad_y.pow(2) + 1e-8)

        # 把红外复制 3 通道再融合
        ir_3ch = ir.repeat(1, 3, 1, 1)
        fused = ir_3ch + grad_mag
        return torch.clamp(fused, 0, 1)
